# encoding: utf-8
import os
import torch
import torchvision.transforms as transforms

tensor_to_pil = transforms.ToPILImage()


def tensor_to_img(tensor_data: torch.Tensor, save_path: str):
    image = tensor_to_pil(tensor_data)
    image.save(save_path)


def batch_tensor_to_img(epoch, out):
    if not os.path.exists("./imgsdemo"):
        os.mkdir("./imgsdemo")
    for idx, tensor in enumerate(out):
        tensor_to_img(tensor, f"./imgsdemo/{epoch}_{idx}.jpg")


def remove_files(file_path, file_type="*"):
    if os.path.exists(file_path) and os.path.isdir(file_path):
        files = os.listdir(file_path)
        for file_name in files:
            os.remove(os.path.join(file_path, file_name))


